import pandas as pd
import torch
from matplotlib import pyplot as plt
import matplotlib.dates as mdates
from model import MyModel  # 导入你的模型类
from dataset import MyDataset  # 导入你的数据集类
from matplotlib.font_manager import FontProperties
import config

opt = config.get_options()
learning_rate = opt.lr
input_size = opt.input_size
hidden_size = opt.hidden_size
dropout_rate = opt.dropout_rate


# 1. 加载训练好的模型
checkpoint_path = r''
model = MyModel.load_from_checkpoint(checkpoint_path,
                                     learning_rate=learning_rate,
                                     input_size=input_size,
                                     hidden_size=hidden_size,
                                     dropout_rate=dropout_rate)
# 2. 读取新数据集
data_path = "./dataset/alpha_charac.xlsx"
new_dataset = MyDataset(data_path)

# 3. 使用模型进行预测
predictions = []
for features, _ in new_dataset:
    # 注意：如果你的模型需要额外的预处理，比如归一化，你需要在这里进行
    features_mean = torch.tensor(new_dataset.features_mean,
                                 dtype=torch.float32)  # Ensure it's a tensor and has float32
    features_std = torch.tensor(new_dataset.features_std,
                                dtype=torch.float32)  # Ensure it's a tensor and has float32
    features = (features - features_mean) / features_std
    with torch.no_grad():
        model.eval()  # 设置模型为评估模式
        output = model(features.unsqueeze(0))  # 对单个样本进行预测
        predictions.append(output.item())

# 4. 将预测结果与原始数据集合并
df = pd.read_excel(data_path)
df['predict_ex_ret'] = predictions

# 5. 使用 calculate_y 函数处理数据
def calculate_y(group):
    group_small = pd.qcut(group['predict_ex_ret'], q=10, precision=10, duplicates='drop')
    group['predict_ex_ret'] = group['predict_ex_ret'].astype(float)
    group['group_sum'] = group.groupby(group_small, observed=False)['predict_ex_ret'].transform('sum')
    group['weight'] = group['predict_ex_ret'] / group['group_sum']
    group['pre_we_y'] = group.groupby(group_small, observed=False)['alpha'].transform(
        lambda x: (x * group.loc[x.index, 'weight']).sum())
    group['equ_we_y'] = group.groupby(group_small, observed=False)['alpha'].transform(lambda x: x.mean())
    group['decile'] = group_small.cat.codes + 1
    return group

# 调用 calculate_y 函数处理数据
df = df.groupby('TradingMonth').apply(calculate_y)
df.reset_index(drop=True, inplace=True)

# 6. 将处理后的数据保存到 Excel 文件中
output_file_path = r"./figure_data\15_data.xlsx"
df.to_excel(output_file_path, index=False)

# 7. 画图
font = FontProperties(fname=r'C:\Windows\Fonts\simhei.ttf', size=12)
# Drop rows where 'TradingMonth' and 'decile' are equal
df = df.drop_duplicates(subset=['TradingMonth', 'decile'])

# Reverse the 'decile' values
df['decile'] = 11 - df['decile']

# Sort the DataFrame by 'TradingMonth' and 'decile'
df = df.sort_values(by=['TradingMonth', 'decile'])

df['TradingMonth'] = pd.to_datetime(df['TradingMonth'])

# Create a figure and axis
fig, ax = plt.subplots(figsize=(10, 6))

# Group by 'decile' column and iterate over the groups
for decile, group in df.groupby('decile'):
    # Calculate cumulative sum for pre_we_y
    group['cum_pre_we_y'] = group['pre_we_y'].cumsum()

    # Plot the line for each group with different color
    ax.plot(group['TradingMonth'], group['cum_pre_we_y'], label=f'Decile {decile}')

# Set x-axis format to show every year
ax.xaxis.set_major_locator(mdates.YearLocator(1))
ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y'))

# Set y-axis interval to -1.5 to 2.5 with 0.2 interval
ax.set_ylim(-2.5, 3.5)
ax.yaxis.set_major_locator(plt.MultipleLocator(0.2))

# Add labels and legend
ax.set_xlabel('Trading Month', fontproperties=font)
ax.set_ylabel('Cumulative pre_we_y', fontproperties=font)
ax.legend()

# Add interactive cursor
plt.gca().format_coord = lambda x, y: f'{x}, {y}'

# Show the plot
plt.show()

# Repeat the same process for equ_we_y
fig, ax = plt.subplots(figsize=(10, 6))

for decile, group in df.groupby('decile'):
    group['cum_equ_we_y'] = group['equ_we_y'].cumsum()
    ax.plot(group['TradingMonth'], group['cum_equ_we_y'], label=f'Decile {decile}')

ax.xaxis.set_major_locator(mdates.YearLocator(1))
ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y'))

# Set y-axis interval to -1.5 to 2.5 with 0.2 interval
ax.set_ylim(-2.5, 3.5)
ax.yaxis.set_major_locator(plt.MultipleLocator(0.2))

ax.set_xlabel('Trading Month', fontproperties=font)
ax.set_ylabel('Cumulative equ_we_y', fontproperties=font)
ax.legend()

# Add interactive cursor
plt.gca().format_coord = lambda x, y: f'{x}, {y}'

plt.show()

# Repeat the same process for real_return
fig, ax = plt.subplots(figsize=(10, 6))

for decile, group in df.groupby('decile'):
    group['cum_equ_we_y'] = group['AdjustedNAVGrowth'].cumsum()
    ax.plot(group['TradingMonth'], group['cum_equ_we_y'], label=f'Decile {decile}')

ax.xaxis.set_major_locator(mdates.YearLocator(1))
ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y'))

# Set y-axis interval to -1.5 to 2.5 with 0.2 interval
ax.set_ylim(-0.5, 5)
ax.yaxis.set_major_locator(plt.MultipleLocator(0.2))

ax.set_xlabel('Trading Month', fontproperties=font)
ax.set_ylabel('Cumulative real_return', fontproperties=font)
ax.legend()

# Add interactive cursor
plt.gca().format_coord = lambda x, y: f'{x}, {y}'

plt.show()